import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 获取股票数据
ticker = "588080.SS"
data = yf.download(ticker, start="2024-01-01", end="2024-05-21")

# 定义短期和长期窗口
short_window = 40
long_window = 100

# 计算移动平均线
data['Short_MA'] = data['Close'].rolling(window=short_window, min_periods=1).mean()
data['Long_MA'] = data['Close'].rolling(window=long_window, min_periods=1).mean()

# 生成交易信号
data['Signal'] = 0
data['Signal'][short_window:] = np.where(data['Short_MA'][short_window:] > data['Long_MA'][short_window:], 1, 0)
data['Position'] = data['Signal'].diff()

# 绘制股票价格和移动平均线
plt.figure(figsize=(12,8))
plt.plot(data['Close'], label='Close Price')
plt.plot(data['Short_MA'], label='40-Day MA')
plt.plot(data['Long_MA'], label='100-Day MA')

# 绘制买卖信号
plt.plot(data[data['Position'] == 1].index, data['Short_MA'][data['Position'] == 1], '^', markersize=10, color='g', label='Buy Signal')
plt.plot(data[data['Position'] == -1].index, data['Short_MA'][data['Position'] == -1], 'v', markersize=10, color='r', label='Sell Signal')

plt.title('AAPL Price with Buy and Sell Signals')
plt.xlabel('Date')
plt.ylabel('Price')
plt.legend()
plt.grid()
plt.show()
